"""该代码的主要目标是渲染PDF的指定页面，并在页面上可视化标注不同类别的段落区域（如标题、图片、表格等）。整体思维逻辑可分为以下几个部分：
  输入：PDF文件路径、文档段落信息列表（doc_list）、目标页码。
  输出：Matplotlib渲染的图像，标注不同类别的段落区域。
关键步骤：
    1.加载PDF页面：使用fitz（PyMuPDF）打开PDF并提取指定页面。
    2.提取目标段落的坐标信息：从doc_list中筛选出目标页面的段落元数据。
    3.渲染PDF为图像：将PDF页面转为位图（pixmap），再转换为PIL图像。
    4.绘制段落框：根据段落坐标和类别，用不同颜色的多边形框标注区域。
    5.添加图例：说明不同颜色对应的段落类别
"""

import fitz  # PyMuPDF库，用于处理PDF文件
import matplotlib.patches as patches  # 用于在图像上绘制多边形
import matplotlib.pyplot as plt  # Matplotlib库，用于绘图
from PIL import Image  # 用于图像处理

def render_pdf_page(file_path, doc_list, page_number):
    """
    渲染指定PDF页面并绘制段落分类框。

    参数：
    - file_path: str，PDF文件路径。
    - doc_list: list，包含段落信息的文档列表，每个元素是一个字典，包含段落元数据。
    - page_number: int，要渲染的页面号（从1开始计数）。
    """
    print("渲染指定PDF页面并绘制段落分类框开始...\n")
    # 打开PDF文件并加载指定页面
    pdf_page = fitz.open(file_path).load_page(page_number - 1)
    segments = [doc.metadata for doc in doc_list if doc.metadata.get("page_number") == page_number]

    # 将PDF页面渲染为位图图像
    pix = pdf_page.get_pixmap()
    pil_image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)

    # 创建绘图环境
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(pil_image)

    # 定义类别颜色映射
    category_to_color = {"Title": "orchid", "Image": "forestgreen", "Table": "tomato"}
    categories = set()

    # 绘制段落标注框
    for segment in segments:
        # 坐标缩放（PDF坐标 → 图像像素坐标）
        points = segment["coordinates"]["points"]
        layout_width = segment["coordinates"]["layout_width"]
        layout_height = segment["coordinates"]["layout_height"]
        scaled_points = [
            (x * pix.width / layout_width, y * pix.height / layout_height) for x, y in points
        ]
        # 绘制多边形框
        box_color = category_to_color.get(segment["category"], "deepskyblue")
        categories.add(segment["category"])
        rect = patches.Polygon(scaled_points, linewidth=1, edgecolor=box_color, facecolor="none")
        ax.add_patch(rect)

    # 添加图例
    legend_handles = [patches.Patch(color="deepskyblue", label="Text")]
    for category, color in category_to_color.items():
        if category in categories:
            legend_handles.append(patches.Patch(color=color, label=category))
    ax.axis("off")
    ax.legend(handles=legend_handles, loc="upper right")
    plt.tight_layout()
    plt.savefig("D:/tmp/渲染指定PDF页面并绘制段落分类框.png")  # 保存渲染结果
    print("渲染指定PDF页面并绘制段落分类框结束...\n")

file_path = ("D:/ideaSpace/rag-in-action-master/90-文档-Data/黑悟空/黑神话悟空.pdf")
from langchain_unstructured import UnstructuredLoader
loader = UnstructuredLoader(
    file_path=file_path,  # PDF文件路径
    strategy="hi_res",    # 使用高分辨率策略进行文档处理
    # partition_via_api=True,  # 通过API进行文档分块
    # coordinates=True,     # 提取文本坐标信息
)
docs = []

# lazy_load() 是一种延迟加载方法
# 它不会一次性将所有文档加载到内存中，而是在需要时才逐个加载文档
# 这对于处理大型PDF文件时可以节省内存使用
for doc in loader.lazy_load():
    docs.append(doc)

print("渲染PDF页面版式\n")
render_pdf_page(file_path,docs, 1)

"""关键逻辑细节
  坐标转换：
    PDF中的段落坐标（layout_width和layout_height）需要缩放到渲染图像的像素尺寸（pix.width和pix.height），确保框的位置和大小正确。
  类别颜色映射：
    默认类别（如文本）用蓝色（deepskyblue），其他类别（标题、图片、表格）通过字典category_to_color映射颜色。
  懒加载文档：
    在调用render_pdf_page前，使用UnstructuredLoader.lazy_load()逐块加载PDF内容，避免内存爆炸。
"""